import json
import os
import sys
import time

import cv2
import numpy as np

from camera.camera import Camera
from robots.piper import PiperArm

import sys

sys.path.append(os.path.dirname(__file__))

class PixelCalibrator:
    def __init__(self, camera: Camera, arm: PiperArm):
        self.image = None
        self.original_image = None
        self.window_name = "Pixel calibration"
        self.scale = 1.0
        self.offset_x = 0
        self.offset_y = 0
        self.dragging = False
        self.last_x, self.last_y = 0, 0
        self.pixels = []
        self.camera_xyz = []
        self.effector_xyz = []
        self.colors = [
            (0, 0, 255),  # 红色
            (0, 255, 0),  # 绿色
            (255, 0, 0),  # 蓝色
            (0, 255, 255),  # 黄色
            (255, 0, 255),  # 紫色
            (255, 255, 0),  # 青色
        ]
        self.time_id = round(time.time())

        self.arm = arm
        self.camera = camera

        # self.restore_points()

    def capture_camera_image(self):
        return self.camera.capture_current_frame()

    def camera_pixel_to_3d(self, x, y):
        return self.camera.get_3d_point_from_pixel(x, y)

    def load_image(self):
        """加载图像"""

        self.original_image = self.capture_camera_image()
        if self.original_image is None:
            print(f"Error: No image captured!")
            return False
        self.image = self.original_image.copy()
        # self.scale = 1.0
        # self.offset_x = 0
        # self.offset_y = 0
        # self.pixels = []
        return True

    def display_coordinates(self, event, x, y, flags, param):
        """鼠标事件回调函数"""
        # 将显示坐标转换为原始图像坐标
        original_x = int((x + self.offset_x) / self.scale)
        original_y = int((y + self.offset_y) / self.scale)

        # 确保坐标在有效范围内
        if (
            original_x < 0
            or original_x >= self.original_image.shape[1]
            or original_y < 0
            or original_y >= self.original_image.shape[0]
        ):
            return

        if event == cv2.EVENT_LBUTTONDOWN:
            # 左键点击：添加标记点
            pixel = [original_x, original_y]
            c_xyz = self.camera_pixel_to_3d(original_x, original_y)
            self.arm.step_ik()  # update simulator status
            end_xyz = self.arm.ik.get_effector_xeuler().tolist()[:3]  # in meters
            # end_xyz = self.arm.get_effector_status()["xyz"].tolist()[:3]

            if c_xyz is not None:
                self.pixels.append(pixel)
                self.camera_xyz.append(c_xyz.tolist())
                self.effector_xyz.append(end_xyz)

            print(f"Add pixel=({original_x}, {original_y}), " f"color={self.original_image[original_y, original_x]}")
            print(f"Camera_xyz={c_xyz}")
            print(f"End_xyz={end_xyz}")

        elif event == cv2.EVENT_RBUTTONDOWN:
            # 右键点击：删除最近添加的点
            if self.pixels:
                removed = self.pixels.pop()
                self.camera_xyz.pop()
                self.effector_xyz.pop()
                print(f"del: pixel=({removed[0]}, {removed[1]})")

        elif event == cv2.EVENT_MBUTTONDOWN:
            # 中键点击：开始拖动
            self.dragging = True
            self.last_x, self.last_y = x, y

        elif event == cv2.EVENT_MOUSEMOVE:
            # 鼠标移动
            if self.dragging:
                # 拖动图像
                dx = x - self.last_x
                dy = y - self.last_y
                self.offset_x += dx
                self.offset_y += dy
                self.last_x, self.last_y = x, y

        elif event == cv2.EVENT_MBUTTONUP:
            # 中键释放：结束拖动
            self.dragging = False

        # 显示当前鼠标位置的坐标
        cv2.displayOverlay(self.window_name, f"Pixel: ({original_x}, {original_y})", 100)

    def run(self):
        """运行标定工具"""

        # image_path = self.camera.save_current_image()
        if not self.load_image():
            return
        img_path, _ = self.camera.save_current_image()
        cv2.imread(img_path)
        cv2.namedWindow(self.window_name, cv2.WINDOW_NORMAL)
        cv2.setMouseCallback(self.window_name, self.display_coordinates)

        while True:
            # 创建当前视图的副本
            display = self.capture_camera_image()

            # 绘制所有标记点
            for i, (px, py) in enumerate(self.pixels):
                color = self.colors[i % len(self.colors)]
                cv2.circle(display, (px, py), 2, color, -1)
                # cv2.putText(display, f"P{i+1}({px}, {py})",
                #            (px+10, py-10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 1)

            # 应用缩放和偏移
            h, w = display.shape[:2]
            view_x1 = max(0, int(self.offset_x / self.scale))
            view_y1 = max(0, int(self.offset_y / self.scale))
            view_x2 = min(w, int((self.offset_x + w) / self.scale))
            view_y2 = min(h, int((self.offset_y + h) / self.scale))

            view = display[view_y1:view_y2, view_x1:view_x2]
            view = cv2.resize(view, (w, h))

            # 显示图像
            cv2.imshow(self.window_name, view)

            # 处理键盘事件
            key = cv2.waitKey(1) & 0xFF

            if key == 27:  # ESC键退出
                break
            elif key == ord("+") or key == ord("="):  # 放大
                self.scale *= 1.1
            elif key == ord("-"):  # 缩小
                self.scale /= 1.1
            elif key == ord("r"):  # 重置视图
                self.scale = 1.0
                self.offset_x = 0
                self.offset_y = 0
                print(self.pixels)
            elif key == ord("s"):  # Save the points
                self.save_points()
            elif key == ord("c"):  # Compute the transformation matrix
                self.compute_transformation()
        cv2.destroyAllWindows()

    def restore_points(self, saved_file: str = "cali_data.json"):

        f_name = f"camera/{self.camera.model}/{self.arm.can}/{saved_file}"

        with open(f_name, "r") as f:
            data = json.load(f)

        self.pixels = data["pixels"]
        self.camera_xyz = data["camera_xyz"]
        self.effector_xyz = data["effector_xyz"]
        print(f"restore {len(self.pixels)} cali_data.json from {f_name}")

    def save_points(self):
        """保存标记点到文件"""
        if not self.pixels:
            print("没有标记点可保存")
            return
        data = {"pixels": self.pixels, "camera_xyz": self.camera_xyz, "effector_xyz": self.effector_xyz}

        f_dir = f"camera/{self.camera.model}/{self.arm.can}"
        if not os.path.exists(f_dir):
            os.mkdir(f_dir)
        f_name = os.path.join(f_dir, f"cali_data_{self.time_id}.json")

        with open(f_name, "w") as f:
            json.dump(data, f)

        print(f"Save {len(self.pixels)} points to {f_name}")

    def compute_transformation(self):
        num_data = len(self.camera_xyz)
        print(f"Compute transformation matrix from {num_data} data...")
        A_hom = np.hstack([np.array(self.camera_xyz), np.ones((num_data, 1))])
        M, _, _, _ = np.linalg.lstsq(A_hom, np.array(self.effector_xyz), rcond=None)

        f_dir = f"camera/{self.camera.model}/{self.arm.can}"
        if not os.path.exists(f_dir):
            os.mkdir(f_dir)
        f_name = os.path.join(f_dir, f"cali_matrix_{self.time_id}.json")

        with open(f_name, "w") as f:
            json.dump(M.tolist(), f)

        print(f"Saved transformation matrix to {f_name}")


if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Use: python pixel_calibration.py <can>")
        sys.exit(1)

    calibrator = PixelCalibrator(
        Camera(),
        PiperArm(sys.argv[1]),
    )
    calibrator.run()

    # python calibration_tool.py can0 130
